#!/usr/bin/env python3

import rospy
import math
import time
from geometry_msgs.msg import Twist
from turtlesim.msg import Pose
from turtlesim.srv import SetPen, SetPenRequest, TeleportAbsolute, TeleportAbsoluteRequest
from std_srvs.srv import Empty

class TurtleNavigator:
    def __init__(self):
        # 初始化ROS节点和通信
        rospy.init_node('turtle_navigator', anonymous=True)
        
        self.velocity_publisher = rospy.Publisher('/turtle1/cmd_vel', Twist, queue_size=10)
        self.pose_subscriber = rospy.Subscriber('/turtle1/pose', Pose, self.pose_callback)
        
        self.setup_services()  # 设置服务连接
        
        self.pose = Pose()
        self.rate = rospy.Rate(20)
        self.stop_movement = False
        
        
        self.set_initial_background()
        
        self.initialize_turtle()  # 初始化海龟位置
        
    def set_initial_background(self):
        """在程序开始时设置初始背景颜色"""
        rospy.loginfo("设置初始背景颜色...")
        self.change_background_color(50, 50, 50)  # 深灰色背景
        
    def setup_services(self):
        # 连接ROS服务
        try:
            rospy.wait_for_service('/turtle1/set_pen', timeout=10)
            rospy.wait_for_service('/turtle1/teleport_absolute', timeout=10)
            rospy.wait_for_service('/clear', timeout=10)
            
            self.set_pen = rospy.ServiceProxy('/turtle1/set_pen', SetPen)
            self.teleport = rospy.ServiceProxy('/turtle1/teleport_absolute', TeleportAbsolute)
            self.clear = rospy.ServiceProxy('/clear', Empty)
            
        except rospy.ROSException as e:
            rospy.logerr("服务连接失败: %s", e)
            raise
    
    def set_pen_color(self, r, g, b, width=2, off=0):
        # 设置画笔颜色和状态
        try:
            pen_req = SetPenRequest()
            pen_req.r = r
            pen_req.g = g
            pen_req.b = b
            pen_req.width = width
            pen_req.off = off
            self.set_pen(pen_req)
        except rospy.ServiceException as e:
            rospy.logerr("画笔设置失败: %s", e)
    
    def teleport_without_trace(self, x, y, theta):
        # 传送海龟
        try:
            self.set_pen_color(0, 0, 0, 1, 1)  # 抬起画笔
            
            req = TeleportAbsoluteRequest()
            req.x = x
            req.y = y
            req.theta = theta
            self.teleport(req)
            
            rospy.sleep(0.1)
            self.set_pen_color(255, 255, 255, 2, 0)  # 放下画笔
            
        except rospy.ServiceException as e:
            rospy.logerr("传送失败: %s", e)
    
    def initialize_turtle(self):
        # 初始化海龟位置
        try:

            self.teleport_without_trace(1.0, 5.5, 0.0)  # 传送到起始点
            rospy.sleep(0.5)
        except rospy.ServiceException as e:
            rospy.logerr("初始化失败: %s", e)
    
    def pose_callback(self, pose_message):
        # 更新海龟位置信息
        self.pose = pose_message
        
    def move_to_point(self, target_x, target_y, speed=2.0, tolerance=0.03):
        # 移动到目标点
        if rospy.is_shutdown() or self.stop_movement:
            return False
            
        vel_msg = Twist()
        success = False
        
        start_time = rospy.Time.now()
        timeout = rospy.Duration(15)
        
        while not rospy.is_shutdown() and not self.stop_movement:
            if rospy.Time.now() - start_time > timeout:
                break
                
            dx = target_x - self.pose.x
            dy = target_y - self.pose.y
            distance = math.sqrt(dx**2 + dy**2)
            
            if distance < tolerance:  # 到达目标点
                success = True
                break
                
            target_angle = math.atan2(dy, dx)
            angle_error = target_angle - self.pose.theta
            angle_error = math.atan2(math.sin(angle_error), math.cos(angle_error))
                
            # 角度控制
            if abs(angle_error) > 0.05:
                vel_msg.angular.z = 3.0 * angle_error
                vel_msg.angular.z = max(-4.0, min(4.0, vel_msg.angular.z))
            else:
                vel_msg.angular.z = 0
                
            # 速度控制
            if abs(angle_error) < 0.3:
                linear_speed = min(speed, distance * 1.5)
                vel_msg.linear.x = min(linear_speed, 3.0)
            else:
                vel_msg.linear.x = 0
                
            self.velocity_publisher.publish(vel_msg)
            self.rate.sleep()
            
        # 停止运动
        vel_msg.linear.x = 0
        vel_msg.angular.z = 0
        self.velocity_publisher.publish(vel_msg)
        
        return success
        
    def change_background_color(self, r, g, b):
        # 改变背景颜色
        try:
            rospy.set_param('/turtlesim/background_r', r)
            rospy.set_param('/turtlesim/background_g', g) 
            rospy.set_param('/turtlesim/background_b', b)
            
            rospy.sleep(0.5)
            self.clear()  # 清除屏幕应用新颜色
            
        except Exception as e:
            rospy.logerr(f"背景颜色设置失败: {e}")
    
    def countdown(self, seconds, message="暂停"):
        # 倒计时功能
        rospy.loginfo(f"{message}，暂停{seconds}秒...")
        for i in range(seconds, 0, -1):
            if rospy.is_shutdown() or self.stop_movement:
                return False
            rospy.loginfo(f"倒计时: {i}秒")
            rospy.sleep(1)
        return True
    
    def reset_turtle_for_next_pattern(self, pattern_number):
        # 重置海龟位置准备下一个图案
        try:
            if not self.countdown(5, f"图案{pattern_number}完成"):
                return
                
            # 改变背景色
            if pattern_number == 1:
                self.change_background_color(0, 100, 0)  # 绿色
            elif pattern_number == 2:
                self.change_background_color(0, 0, 100)  # 蓝色
            
            # 传送到新位置
            if pattern_number == 1:
                self.teleport_without_trace(9.0, 9.0, math.pi)
            elif pattern_number == 2:
                self.teleport_without_trace(2.0, 2.0, math.pi/2)
                
        except rospy.ServiceException as e:
            rospy.logerr("重置失败: %s", e)
        
    def draw_circle_with_s_curve(self, radius=2.0):
        # 绘制圆形带S曲线
        if rospy.is_shutdown() or self.stop_movement:
            return
            
        rospy.loginfo("开始绘制圆形带S曲线...")
        
        # 移动到起始位置
        start_x = 5.5 - radius
        start_y = 5.5
        self.teleport_without_trace(start_x, start_y, 0)
        
        # 绘制外圆
        rospy.loginfo("绘制外圆...")
        center_x, center_y = 5.5, 5.5
        num_points = 60
        
        # 移动到第一个点
        theta0 = 0
        x0 = center_x + radius * math.cos(theta0)
        y0 = center_y + radius * math.sin(theta0)
        self.teleport_without_trace(x0, y0, theta0 + math.pi/2)
        
        # 绘制圆形
        for i in range(1, num_points + 1):
            if rospy.is_shutdown() or self.stop_movement:
                return
            theta = 2 * math.pi * i / num_points
            x = center_x + radius * math.cos(theta)
            y = center_y + radius * math.sin(theta)
            if not self.move_to_point(x, y, speed=2.5):
                return
        
        # 停止并抬起画笔
        vel_msg = Twist()
        vel_msg.linear.x = 0
        vel_msg.angular.z = 0
        self.velocity_publisher.publish(vel_msg)
        rospy.sleep(0.1)
        self.set_pen_color(0, 0, 0, 1, 1)
        rospy.sleep(0.1)
        
        # 绘制S曲线
        rospy.loginfo("绘制S形曲线...")
        s_radius = radius / 2
        s_start_x = center_x - 2 * s_radius
        s_start_y = center_y
        
        # 传送到S曲线起点
        self.teleport_without_trace(s_start_x, s_start_y, math.pi/2)
        
        # 绘制第一部分半圆
        num_points_s = 30
        for i in range(0, num_points_s + 1):
            if rospy.is_shutdown() or self.stop_movement:
                return
            theta = math.pi - (math.pi * i / num_points_s)
            x = (center_x - s_radius) + s_radius * math.cos(theta)
            y = center_y + s_radius * math.sin(theta)
            if not self.move_to_point(x, y, speed=2.5):
                return
        
        # 停止并抬起画笔
        vel_msg.linear.x = 0
        vel_msg.angular.z = 0
        self.velocity_publisher.publish(vel_msg)
        rospy.sleep(0.1)
        self.set_pen_color(0, 0, 0, 1, 1)
        rospy.sleep(0.1)
        
        # 传送到第二部分起点
        self.teleport_without_trace(center_x, center_y, math.pi/2)
        
        # 绘制第二部分半圆
        for i in range(0, num_points_s + 1):
            if rospy.is_shutdown() or self.stop_movement:
                return
            theta = math.pi * i / num_points_s
            x = (center_x + s_radius) - s_radius * math.cos(theta)
            y = center_y - s_radius * math.sin(theta)
            if not self.move_to_point(x, y, speed=2.5):
                return
        
        rospy.loginfo("圆形带S曲线绘制完成!")
        
    def draw_gear_square_wave(self, R=2.0, h=0.4, num_teeth=8, center_x=7.0, center_y=7.0):
        # 绘制方波齿轮
        if rospy.is_shutdown() or self.stop_movement:
            return
            
        rospy.loginfo("开始绘制方波齿轮...")
        
        num_points = num_teeth * 20
        
        # 计算起始位置和方向
        theta0 = 0
        r0 = self._get_square_wave_radius(theta0, R, h, num_teeth)
        x0 = center_x + r0 * math.cos(theta0)
        y0 = center_y + r0 * math.sin(theta0)
        
        theta_small = 0.01
        r_small = self._get_square_wave_radius(theta_small, R, h, num_teeth)
        x_small = center_x + r_small * math.cos(theta_small)
        y_small = center_y + r_small * math.sin(theta_small)
        dx = x_small - x0
        dy = y_small - y0
        initial_angle = math.atan2(dy, dx)
        
        # 传送到起始点
        self.teleport_without_trace(x0, y0, initial_angle)
        
        # 绘制齿轮
        for i in range(1, num_points + 1):
            if rospy.is_shutdown() or self.stop_movement:
                return
                
            theta = 2 * math.pi * i / num_points
            r = self._get_square_wave_radius(theta, R, h, num_teeth)
            x = center_x + r * math.cos(theta)
            y = center_y + r * math.sin(theta)
            
            if not self.move_to_point(x, y, speed=1.5):
                return
            
        rospy.loginfo("方波齿轮绘制完成!")
    
    def _get_square_wave_radius(self, theta, R, h, num_teeth):
        # 计算方波齿轮半径
        theta_norm = theta % (2 * math.pi)
        tooth_period = 2 * math.pi / num_teeth
        tooth_phase = theta_norm % tooth_period
        
        # 方波：前半个周期为齿，后半个周期为谷
        if tooth_phase < tooth_period / 2:
            return R + h  # 齿
        else:
            return R      # 谷
        
    def run(self):
        # 主程序：依次绘制两个图案
        try:
            rospy.loginfo("开始绘制图案...")
            
            # 绘制第一个图案
            self.draw_circle_with_s_curve(radius=2.0)
            
            if rospy.is_shutdown() or self.stop_movement:
                return
                
            # 重置并改变背景
            self.reset_turtle_for_next_pattern(1)
            
            if rospy.is_shutdown() or self.stop_movement:
                return
                
            # 绘制第二个图案
            self.draw_gear_square_wave(R=2.0, h=0.4, num_teeth=8, center_x=7.0, center_y=7.0)
            
            if rospy.is_shutdown() or self.stop_movement:
                return
                
            # 最终重置
            self.reset_turtle_for_next_pattern(2)
            
            rospy.loginfo("所有图形绘制完成!")
            
        except Exception as e:
            rospy.logerr("程序执行出错: %s", e)

if __name__ == '__main__':
    try:
        navigator = TurtleNavigator()
        navigator.run()
    except Exception as e:
        rospy.logerr("程序启动失败: %s", e)
